package com.xyzboom.compiler.fa

private typealias NFANodeD<T, K> = DFANode<MutableSet<T>, K>

private inline fun <T, K> epsilonAction(
    dfaState: MutableSet<T>, stateMap: HashMap<T, HashMap<FAEdge<K>?, HashSet<T>>>
) {
    val needVisit = ArrayDeque(dfaState)
    val visited = HashSet<T>()
    val it = needVisit.iterator()
    while (it.hasNext()) {
        val nfaState = it.next()
        it.remove()
        if (visited.contains(nfaState)) continue
        visited.add(nfaState)
        val stateTo = stateMap[nfaState] ?: continue
        val stateAddToDfa = stateTo[null] ?: continue
        dfaState.addAll(stateAddToDfa)
        stateAddToDfa.forEach { needVisit.add(it) }
    }
}

private fun <T, K> moveAction(
    dfaState: MutableSet<T>, action: FAEdge<K>, stateMap: HashMap<T, HashMap<FAEdge<K>?, HashSet<T>>>
): DFANode<MutableSet<T>, K> {
    val moveResult = HashSet<T>()
    for (nfaState in dfaState) {
        val stateTo = stateMap[nfaState] ?: continue
        stateTo[action]?.let { moveResult.addAll(it) }
    }
    return DFANode(moveResult)
}

private fun <T, K> NFA<T, K>.getStateMap(): HashMap<T, HashMap<FAEdge<K>?, HashSet<T>>> {
    val nfaStateMap = HashMap<T, HashMap<FAEdge<K>?, HashSet<T>>>()
    val visited = HashSet<T>()
    for (nodeNow in allNodes) {
        if (visited.contains(nodeNow.value)) {
            continue
        }
        for (entryNext in nodeNow.next) {
            nfaStateMap.getOrPut(nodeNow.value) { HashMap() }
                .getOrPut(entryNext.key) { HashSet() }
                .addAll(entryNext.value.map {
                    it.value
                })
        }
        visited.add(nodeNow.value)
    }
    return nfaStateMap
}

fun <T, K> NFA<T, K>.toDFA(): DFA<Int, K> {
    val nfaStateMap = getStateMap()
    val dfaStartState = NFANodeD<T, K>(HashSet())
    val dfaEndStates = HashSet<NFANodeD<T, K>>()
    val dfaNormalStates = HashSet<NFANodeD<T, K>>()
    val dfaAllStates = HashMap<NFANodeD<T, K>, Boolean>()
    dfaStartState.value.add(nodeIn.value)
    epsilonAction(dfaStartState.value, nfaStateMap)
    if (dfaStartState.value.contains(nodeOut.value)) {
        dfaEndStates.add(dfaStartState)
    }
    dfaAllStates[dfaStartState] = false
    while (dfaAllStates.containsValue(false)) {
        var stateNow: NFANodeD<T, K>? = null
        run {
            dfaAllStates.forEach { (k, v) ->
                if (!v) {
                    stateNow = k
                    return@run
                }
            }
        }
        stateNow ?: continue
        val guessAlphabet = HashSet<FAEdge<K>>()
        for (nfaState in stateNow!!.value) {
            nfaStateMap[nfaState]?.keys?.forEach {
                it ?: return@forEach
                guessAlphabet.add(it)
            }
        }
        for (k in guessAlphabet) {
            val moveResult = moveAction(stateNow!!.value, k, nfaStateMap)
            epsilonAction(moveResult.value, nfaStateMap)
            if (moveResult.value.isEmpty()) {
                continue
            }
            if (!dfaAllStates.containsKey(moveResult)) {
                dfaAllStates[moveResult] = false
                if (!moveResult.value.contains(nodeOut.value)) {
                    dfaNormalStates.add(moveResult)
                } else {
                    dfaEndStates.add(moveResult)
                }
            }
            stateNow!!.next[k] = moveResult
        }
        dfaAllStates[stateNow!!] = true
    }
    return DFA(dfaStartState, dfaEndStates, dfaAllStates.keys).intStateDFA().getSimplest().intStateDFA()
}